#define _GNU_SOURCE

#include <stdio.h>
#include <unistd.h>
#include <dlfcn.h>
#include <string.h>
#include <stdlib.h>
#include <stdbool.h>
#include <spawn.h>
#include <regex.h>
#include <libgen.h>
#include <fcntl.h>

#define bool int
#define true 1
#define false 0

int (*real_execve)(const char*, char* const*, char* const*);
int (*real_posix_spawn)(pid_t*, const char*, const posix_spawn_file_actions_t*, const posix_spawnattr_t*, char* const*, char* const*);
int (*real_posix_spawnp)(pid_t*, const char*, const posix_spawn_file_actions_t*, const posix_spawnattr_t*, char* const*, char* const*);

char* so_path = SNIFFER_SO_PATH;
char* fifo_path = SNIFFER_FIFO_PATH;
char* old_ld_preload = NULL;

void try_to_intercept(const char* filepath, const char** argv, const char** envp)
{
    const char* fifo_path = getenv("command_sniffer_fifo");
    if (fifo_path == NULL) {
        fprintf(stderr, "Warning: command_sniffer_fifo not set\n");
        return;
    }
    char cwd[1024];
    getcwd(cwd, sizeof(cwd));
    static char buffer[1024 * 100];
    sprintf(buffer, "working_directory: %s\n", cwd);
    for (int i = 0; argv[i] != NULL; i++) {
        strcat(buffer, "command:");
        strcat(buffer, argv[i]);
        strcat(buffer, "\n");
    }
    strcat(buffer, "\nenviron:\n");
    for (int i = 0; envp[i] != NULL; i++) {
        strcat(buffer, "  ");
        strcat(buffer, envp[i]);
        strcat(buffer, "\n");
    }
    int fifo_fd = open(fifo_path, O_WRONLY);
    if (fifo_fd < 0) {
        fprintf(stderr, "Error opening fifo file: %s\n", fifo_path);
        exit(1);
    }
    int fl = flock(fifo_fd, LOCK_EX);
    if (fl == -1) {
        fprintf(stderr, "Error locking fifo file: %s\n", fifo_path);
        exit(1);
    }
    if (write(fifo_fd, buffer, strlen(buffer) + 1) == -1) {
        fprintf(stderr, "Error writing to fifo file: %s\n", fifo_path);
        exit(1);
    }
    flock(fifo_fd, LOCK_UN);
    close(fifo_fd);
}

static bool is_sniffer_being_loaded(char** envp)
{
    for (int i = 0; envp[i] != NULL; i++) {
        if (strncmp(envp[i], "LD_PRELOAD=", 11) == 0) {
            char* so_path = strstr(envp[i], so_path);
            if (so_path != NULL) {
                return true;
            } else {
                old_ld_preload = envp[i] + 11;
                return false;
            }
        }
    }
    return false;
}

static char** refresh_envp(char** envp)
{
    if (is_sniffer_being_loaded(envp)) {
        return envp;
    }
    int envp_size = 0;
    for (int i = 0; envp[i] != NULL; i++) {
        envp_size++;
    }
    char** new_envp = (char**)malloc((envp_size + 3) * sizeof(char*));
    static char so_path_buffer[1024] = { 0 };
    static char fifo_path_buffer[1024] = { 0 };
    sprintf(so_path_buffer, "LD_PRELOAD=%s", so_path);
    if (old_ld_preload) {
        strcat(so_path_buffer, old_ld_preload);
    }
    sprintf(fifo_path_buffer, "command_sniffer_fifo=%s", fifo_path);
    memcpy(new_envp, envp, sizeof(char*) * envp_size);
    new_envp[envp_size] = so_path_buffer;
    new_envp[envp_size + 1] = fifo_path_buffer;
    new_envp[envp_size + 2] = NULL;
    return new_envp;
}

int execve(const char* filepath, char* const argv[], char* const envp[])
{
    if (real_execve == NULL) {
        real_execve = dlsym(RTLD_NEXT, "execve");
    }
    try_to_intercept(filepath, argv, envp);
    char** new_envp = refresh_envp(envp);
    if (new_envp != envp) {
        int ret = real_execve(filepath, argv, new_envp);
        free(new_envp);
        return ret;
    }
    return real_execve(filepath, argv, envp);
}

int posix_spawn(pid_t* restrict pid, const char* restrict filepath,
                const posix_spawn_file_actions_t* restrict file_actions,
                const posix_spawnattr_t* restrict attrp,
                char* const argv[restrict],
                char* const envp[restrict])
{
    if (real_posix_spawn == NULL) {
        real_posix_spawn = dlsym(RTLD_NEXT, "posix_spawn");
    }
    try_to_intercept(filepath, argv, envp);
    return real_posix_spawn(pid, filepath, file_actions, attrp, argv, envp);
}

int posix_spawnp(pid_t* restrict pid, const char* restrict filepath,
                 const posix_spawn_file_actions_t* restrict file_actions,
                 const posix_spawnattr_t* restrict attrp,
                 char* const argv[restrict],
                 char* const envp[restrict])
{
    if (real_posix_spawnp == NULL) {
        real_posix_spawnp = dlsym(RTLD_NEXT, "posix_spawnp");
    }
    try_to_intercept(filepath, argv, envp);
    return real_posix_spawnp(pid, filepath, file_actions, attrp, argv, envp);
}